import cv2
import numpy as np
import os

def enhance_edge_detection(image):
    """增强边缘检测精度"""
    # 使用多通道多方法边缘检测
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    
    # 多方法边缘检测
    edges_gray = cv2.Canny(gray, 50, 150)
    edges_s = cv2.Canny(s, 30, 100)
    edges_v = cv2.Canny(v, 30, 100)
    
    # 组合边缘
    combined_edges = cv2.bitwise_or(edges_gray, edges_s)
    combined_edges = cv2.bitwise_or(combined_edges, edges_v)
    
    # 形态学优化
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    enhanced_edges = cv2.morphologyEx(combined_edges, cv2.MORPH_CLOSE, kernel)
    enhanced_edges = cv2.dilate(enhanced_edges, kernel, iterations=1)
    
    return enhanced_edges

def split_color_blocks():
    # 读取图像
    img = cv2.imread('./test.jpg')
    if img is None:
        print("Error: Image not found!")
        return
    
    # 创建输出文件夹
    os.makedirs('./split', exist_ok=True)
    
    # 增强边缘检测
    enhanced_edges = enhance_edge_detection(img)
    
    # 查找轮廓并获取层次结构
    contours, hierarchy = cv2.findContours(enhanced_edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    
    # 检查是否找到轮廓
    if not contours:
        print("No contours found!")
        return
    
    # 将层次结构转换为更易处理的格式
    hierarchy = hierarchy[0]  # 获取实际的层次数组
    
    # 创建带透明通道的结果图像
    bgra_img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
    
    # 处理每个轮廓
    block_count = 0
    for i, cnt in enumerate(contours):
        # 跳过内部轮廓（只处理外部轮廓）
        if hierarchy[i][3] != -1:  # 不是最外层轮廓
            continue
            
        # 创建当前色块掩码
        block_mask = np.zeros(enhanced_edges.shape, dtype=np.uint8)
        cv2.drawContours(block_mask, [cnt], -1, 255, -1)
        
        # 处理孔洞 - 查找并扣除子轮廓
        child_idx = hierarchy[i][2]
        while child_idx != -1:
            # 确保索引有效
            if child_idx < len(contours):
                cv2.drawContours(block_mask, [contours[child_idx]], -1, 0, -1)
            # 移动到下一个同级轮廓
            child_idx = hierarchy[child_idx][0]
        
        # 创建当前色块图像
        block_img = bgra_img.copy()
        block_img[:, :, 3] = np.where(block_mask == 255, 255, 0)
        
        # 保存结果
        cv2.imwrite(f'./split/color_block_{block_count:04d}.png', block_img)
        block_count += 1

    print(f"Saved {block_count} color blocks to ./split folder")

if __name__ == "__main__":
    split_color_blocks()